
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

#include "ff_wrap.h"
#include "record.h"

const static int codec_map[REC_CODEC_MAX] = 
{
	AV_CODEC_ID_H264,
	AV_CODEC_ID_MJPEG,
	AV_CODEC_ID_AAC,
	AV_CODEC_ID_PCM_ALAW
};

/******************************************************************
 *   FFmpeg Muxer Wrapper Functions
 *****************************************************************/
/* Add an output stream. */
static AVStream *add_stream(st_ff_context *ff_ctx,enum AVCodecID codec_id, void *param)
{
	AVFormatContext *oc = ff_ctx->fmt_ctx;
	AVCodec *codec;
    AVCodecContext *c;
    AVStream *st;

	struct rsc_record_param *rec_param = (struct rsc_record_param *)param;

    /* find the encoder */
    codec = avcodec_find_encoder(codec_id);
    if (!codec) 
	{
        fprintf(stderr, "Could not find encoder for '%s',try decoder...\n",
                avcodec_get_name(codec_id));
		/** just for getting AVCodec info: type, id */
		codec = avcodec_find_decoder(codec_id);
		if(!codec) return NULL;
    }

	/** codec need: type, id */
    st = avformat_new_stream(oc, codec);
    if (!st) 
	{
        fprintf(stderr, "Could not allocate stream\n");
        return NULL;
    }
    st->id = oc->nb_streams-1;
    c = st->codec;

    switch (codec->type) {
    case AVMEDIA_TYPE_AUDIO:
		ff_ctx->audio_codec = codec;
		ff_ctx->audio_stream_idx = st->id;

		if(rec_param)
		{
			c->sample_fmt  = rec_param->media.audio_sample_fmt;
			c->bit_rate    = rec_param->media.audio_bitrate;
			c->sample_rate = rec_param->media.audio_sample_rate;
			c->channels    = rec_param->media.audio_channels;
			c->channel_layout = av_get_default_channel_layout(c->channels);
		}
		else
		{
			c->sample_fmt  = STREAM_SAMPLE_FMT;
			c->bit_rate    = 64000;
			c->sample_rate = 44100;
			c->channels    = 2;
			c->channel_layout = av_get_default_channel_layout(c->channels);
		}
        break;

    case AVMEDIA_TYPE_VIDEO:
		ff_ctx->video_codec = codec;
		ff_ctx->video_stream_idx = st->id;

        c->codec_id = codec_id;
		if(rec_param)
		{
			c->bit_rate = rec_param->media.video_bitrate;
			/* Resolution must be a multiple of two. */
			c->width    = rec_param->media.video_width;
			c->height   = rec_param->media.video_height;
			/* timebase: This is the fundamental unit of time (in seconds) in terms
			 * of which frame timestamps are represented. For fixed-fps content,
			 * timebase should be 1/framerate and timestamp increments should be
			 * identical to 1. */
			c->time_base.den = rec_param->media.video_framerate;
			c->time_base.num = 1;
			/* emit one intra frame every twelve frames at most */
			c->gop_size      = rec_param->media.video_gop; 
			c->pix_fmt       = STREAM_PIX_FMT;
		}
		else
		{
			c->bit_rate = 400000;
			/* Resolution must be a multiple of two. */
			c->width    = 352;
			c->height   = 288;
			/* timebase: This is the fundamental unit of time (in seconds) in terms
			 * of which frame timestamps are represented. For fixed-fps content,
			 * timebase should be 1/framerate and timestamp increments should be
			 * identical to 1. */
			c->time_base.den = 25;
			c->time_base.num = 1;
			c->gop_size      = 12; /* emit one intra frame every twelve frames at most */
			c->pix_fmt       = STREAM_PIX_FMT;
			if (c->codec_id == AV_CODEC_ID_MPEG2VIDEO) {
				/* just for testing, we also add B frames */
				c->max_b_frames = 2;
			}
			if (c->codec_id == AV_CODEC_ID_MPEG1VIDEO) {
				/* Needed to avoid using macroblocks in which some coeffs overflow.
				 * This does not happen with normal video, it just happens here as
				 * the motion of the chroma plane does not match the luma plane. */
				c->mb_decision = 2;
			}
		}
		break;

    default:
        break;
    }

    /* Some formats want stream headers to be separate. */
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

#if FF_CODEC_OPEN
	int ret = 0;
    /* open the codec */
    ret = avcodec_open2(c, codec, NULL);
    if (ret < 0) 
	{
        fprintf(stderr, "Could not open video codec: %s\n", av_err2str(ret));
        return NULL;
    }
#endif

    return st;
}

int ff_write_audio_frame(st_ff_context *ff_ctx, unsigned char *data,int len)
{
	AVFormatContext *oc = ff_ctx->fmt_ctx;
	AVStream *st = ff_ctx->audio_st; 
    AVPacket pkt = { 0 }; // data and size must be 0;
	int ret = 0;

    av_init_packet(&pkt);

	pkt.flags        |= AV_PKT_FLAG_KEY;
	pkt.stream_index  = st->index;
	pkt.data          = data;
	pkt.size          = len;
	pkt.pts = pkt.dts = AV_NOPTS_VALUE;

    /* Write the compressed frame to the media file. */
    ret = av_interleaved_write_frame(oc, &pkt);
    if (ret != 0) 
	{
        fprintf(stderr, "Error while writing audio frame: %s\n",
                av_err2str(ret));
        return -1;
    }

    ff_ctx->aframe_cnt++;
	return 0;
}

int ff_write_video_frame(st_ff_context *ff_ctx, unsigned char *data,int len)
{
	int ret = 0;

	if(ff_ctx == NULL || data == NULL || len <= 0)
	{
		printf("invalid params!\n");
		return -1;
	}

	AVFormatContext *oc = ff_ctx->fmt_ctx;
	AVStream *st = ff_ctx->video_st;

	AVPacket pkt;
	av_init_packet(&pkt);

	pkt.flags        |= AV_PKT_FLAG_KEY;// keyframe
	pkt.stream_index  = st->index;
	pkt.data          = data;
	pkt.size          = len;
	pkt.pts = pkt.dts = AV_NOPTS_VALUE;

	ret = av_interleaved_write_frame(oc, &pkt);

    if (ret != 0) 
	{
        fprintf(stderr, "Error while writing video frame: %s\n", av_err2str(ret));
        return -1;
    }

    ff_ctx->vframe_cnt++;
	return 0;
}

static void close_audio(st_ff_context *ff_ctx)
{
	if(ff_ctx==NULL)
		return;
	//AVFormatContext *oc = ff_ctx->fmt_ctx;
	AVStream *st = ff_ctx->audio_st;
	if(st==NULL)
		return;

    avcodec_close(st->codec);
}
static void close_video(st_ff_context *ff_ctx)
{
	if(ff_ctx==NULL)
		return;

	//AVFormatContext *oc = ff_ctx->fmt_ctx;
	AVStream *st = ff_ctx->video_st;

	if(st==NULL)
		return;

    avcodec_close(st->codec);
}


int ff_init_muxer(const char *filename, st_ff_context *ff_ctx, void *param)
{
	if(filename == NULL || ff_ctx == NULL)
	{
		printf("invalid params!\n");
		return -1;
	}
	int ret = 0;

	memset(ff_ctx,0,sizeof(st_ff_context));
    /* Initialize libavcodec, and register all codecs and formats. */
    av_register_all();
	
    /* allocate the output media context */
    avformat_alloc_output_context2(&ff_ctx->fmt_ctx, NULL, NULL, filename);
    if (!ff_ctx->fmt_ctx) 
	{
        printf("Could not deduce output format from file extension: using avi.\n");
        avformat_alloc_output_context2(&ff_ctx->fmt_ctx, NULL, "avi", filename);
    }
	if(!ff_ctx->fmt_ctx)
		return -1;

	AVOutputFormat *fmt = ff_ctx->fmt_ctx->oformat;

    /* Add the audio and video streams using the default format codecs
     * and initialize the codecs. */
    ff_ctx->video_st = NULL;
    ff_ctx->audio_st = NULL;

	struct rsc_record_param *rec_param = (struct rsc_record_param *)param;
    if (fmt->video_codec != AV_CODEC_ID_NONE) 
	{
		printf("video codec: %s\n",avcodec_get_name(fmt->video_codec));
		if(rec_param) 
			fmt->video_codec = codec_map[rec_param->media.video_codec];
        ff_ctx->video_st = add_stream(ff_ctx, fmt->video_codec,param);
    }
    if (fmt->audio_codec != AV_CODEC_ID_NONE) 
	{
		printf("audio codec: %s\n",avcodec_get_name(fmt->audio_codec));
		if(rec_param) 
			fmt->audio_codec = codec_map[rec_param->media.audio_codec];
        ff_ctx->audio_st = add_stream(ff_ctx, fmt->audio_codec,param);
    }
	if(!ff_ctx->video_st || !ff_ctx->audio_st)
	{
		printf("add stream failed!\n");
		return -1;
	}

    av_dump_format(ff_ctx->fmt_ctx, 0, filename, 1);

    /* open the output file, if needed */
    if (!(fmt->flags & AVFMT_NOFILE)) 
	{
        ret = avio_open(&ff_ctx->fmt_ctx->pb, filename, AVIO_FLAG_WRITE);
        if (ret < 0) 
		{
            fprintf(stderr, "Could not open '%s': %s\n", filename,
                    av_err2str(ret));
            return -1;
        }
    }

    /* Write the stream header, if any. */
    ret = avformat_write_header(ff_ctx->fmt_ctx, NULL);
    if (ret < 0) 
	{
        fprintf(stderr, "Error occurred when opening output file: %s\n",
                av_err2str(ret));
        return -1;
    }
	return 0;
}

int ff_exit_muxer(st_ff_context *ff_ctx)
{
	if(ff_ctx==NULL)
		return -1;

	AVFormatContext *oc = ff_ctx->fmt_ctx;

    /* Write the trailer, if any. The trailer must be written before you
     * close the CodecContexts open when you wrote the header; otherwise
     * av_write_trailer() may try to use memory that was freed on
     * av_codec_close(). */
    av_write_trailer(oc);

    /* Close each codec. */
	close_video(ff_ctx);
	close_audio(ff_ctx);

    if (!(oc->oformat->flags & AVFMT_NOFILE))
        /* Close the output file. */
        avio_close(oc->pb);

    /* free the stream */
    avformat_free_context(oc);

	memset(ff_ctx,0,sizeof(st_ff_context));
    return 0;
}



/******************************************************************
 *   FFmpeg DeMuxer Wrapper Functions
 *****************************************************************/
static int open_codec_context(st_ff_context *ff_ctx, enum AVMediaType type)
{
    int ret;
	int stream_idx = 0;
	AVFormatContext *fmt_ctx = ff_ctx->fmt_ctx;
    AVStream *st;
    AVCodecContext *dec_ctx = NULL;
    AVCodec *dec = NULL;

    ret = av_find_best_stream(fmt_ctx, type, -1, -1, NULL, 0);
    if (ret < 0) 
	{
        fprintf(stderr, "Could not find %s stream in input file\n",
                av_get_media_type_string(type));
        return ret;
    } 
	else 
	{
        stream_idx = ret;
        st = fmt_ctx->streams[stream_idx];

        /* find decoder for the stream */
        dec_ctx = st->codec;
        dec = avcodec_find_decoder(dec_ctx->codec_id);
        if (!dec) 
		{
            fprintf(stderr, "Failed to find %s codec\n",
                    av_get_media_type_string(type));
            return ret;
        }
#if FF_CODEC_OPEN
        if ((ret = avcodec_open2(dec_ctx, dec, NULL)) < 0) 
		{
            fprintf(stderr, "Failed to open %s codec\n",
                    av_get_media_type_string(type));
            return ret;
        }
#endif
    }
	if(type == AVMEDIA_TYPE_VIDEO)
	{
		ff_ctx->video_stream_idx = stream_idx;
		ff_ctx->video_st = st;
		ff_ctx->video_codec = dec;
	}
	else if(type == AVMEDIA_TYPE_AUDIO)
	{
		ff_ctx->audio_stream_idx = stream_idx;
		ff_ctx->audio_st = st;
		ff_ctx->audio_codec = dec;
	}

    return 0;
}

int ff_init_demuxer(const char *filename, st_ff_context *ff_ctx)
{
	if(filename == NULL || ff_ctx == NULL)
	{
		printf("invalid params!\n");
		return -1;
	}

	//AVCodecContext *video_dec_ctx = NULL, *audio_dec_ctx;
	AVStream *video_stream = NULL, *audio_stream = NULL;

	memset(ff_ctx,0,sizeof(st_ff_context));
    /* Initialize libavcodec, and register all codecs and formats. */
    av_register_all();

    /* open input file, and allocate format context */
    if (avformat_open_input(&ff_ctx->fmt_ctx, filename, NULL, NULL) < 0)
	{
        fprintf(stderr, "Could not open source file %s\n", filename);
        return 1;
    }

    /* retrieve stream information */
    if (avformat_find_stream_info(ff_ctx->fmt_ctx, NULL) < 0) 
	{
        fprintf(stderr, "Could not find stream information\n");
        return 1;
    }

    if (open_codec_context(ff_ctx, AVMEDIA_TYPE_VIDEO) >= 0)
	{
        video_stream = ff_ctx->video_st;
        //video_dec_ctx = video_stream->codec;
    }

    if (open_codec_context(ff_ctx, AVMEDIA_TYPE_AUDIO) >= 0) 
	{
        audio_stream = ff_ctx->audio_st;
        //audio_dec_ctx = audio_stream->codec;
    }

    /* dump input information to stderr */
    av_dump_format(ff_ctx->fmt_ctx, 0, filename, 0);

    if (!audio_stream && !video_stream) 
	{
        fprintf(stderr, "Could not find audio or video stream in the input, aborting\n");
        return -1;
    }

	return 0;
}

int ff_exit_demuxer(st_ff_context *ff_ctx)
{
	if(ff_ctx == NULL)
		return -1;
	AVCodecContext *video_dec_ctx = ff_ctx->video_st->codec;
	AVCodecContext *audio_dec_ctx = ff_ctx->audio_st->codec;

    if (video_dec_ctx)
        avcodec_close(video_dec_ctx);
    if (audio_dec_ctx)
        avcodec_close(audio_dec_ctx);

    avformat_close_input(&ff_ctx->fmt_ctx);
	memset(ff_ctx,0,sizeof(st_ff_context));

	return 0;
}

